// Copyright (C) Kumo inc. and its affiliates.
// Author: Jeff.li lijippy@163.com
// All rights reserved.
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.
//
#include <pollux/functions/prestosql/aggregates/map_union_aggregate.h>
#include <pollux/functions/prestosql/aggregates/map_aggregate_base.h>

namespace kumo::pollux::aggregate::prestosql {
    namespace {
        // See documentation at
        // https://prestodb.io/docs/current/functions/aggregate.html
        template<typename K, typename AccumulatorType = MapAccumulator<K> >
        class MapUnionAggregate : public MapAggregateBase<K, AccumulatorType> {
        public:
            using Base = MapAggregateBase<K, AccumulatorType>;

            explicit MapUnionAggregate(TypePtr resultType) : Base(resultType) {
            }

            bool supportsToIntermediate() const override {
                return true;
            }

            void toIntermediate(
                const SelectivityVector &rows,
                std::vector<VectorPtr> &args,
                VectorPtr &result) const override {
                if (rows.isAllSelected()) {
                    result = args[0];
                } else {
                    auto *pool = Base::allocator_->pool();
                    const auto numRows = rows.size();

                    // Set nulls for rows not present in 'rows'.
                    BufferPtr nulls = allocate_nulls(numRows, pool);
                    memcpy(
                        nulls->asMutable<uint64_t>(),
                        rows.asRange().bits(),
                        bits::nbytes(numRows));

                    BufferPtr indices = allocate_indices(numRows, pool);
                    auto *rawIndices = indices->asMutable<vector_size_t>();
                    std::iota(rawIndices, rawIndices + numRows, 0);
                    result =
                            BaseVector::wrap_in_dictionary(nulls, indices, rows.size(), args[0]);
                }
            }

            void addRawInput(
                char **groups,
                const SelectivityVector &rows,
                const std::vector<VectorPtr> &args,
                bool /*mayPushdown*/) override {
                Base::addMapInputToAccumulator(groups, rows, args, false);
            }

            void addSingleGroupRawInput(
                char *group,
                const SelectivityVector &rows,
                const std::vector<VectorPtr> &args,
                bool /*mayPushdown*/) override {
                Base::addSingleGroupMapInputToAccumulator(group, rows, args, false);
            }
        };

        template<TypeKind Kind>
        std::unique_ptr<exec::Aggregate> createMapUnionAggregateWithCustomCompare(
            const TypePtr &resultType) {
            return std::make_unique<MapUnionAggregate<
                typename TypeTraits<Kind>::NativeType,
                CustomComparisonMapAccumulator<Kind> > >(resultType);
        }
    } // namespace

    void registerMapUnionAggregate(
        const std::string &prefix,
        bool withCompanionFunctions,
        bool overwrite) {
        std::vector<std::shared_ptr<exec::AggregateFunctionSignature> > signatures{
            exec::AggregateFunctionSignatureBuilder()
            .typeVariable("K")
            .typeVariable("V")
            .returnType("map(K,V)")
            .intermediateType("map(K,V)")
            .argumentType("map(K,V)")
            .build()
        };

        auto name = prefix + kMapUnion;
        exec::registerAggregateFunction(
            name,
            std::move(signatures),
            [name](
        core::AggregationNode::Step /*step*/,
        const std::vector<TypePtr> &argTypes,
        const TypePtr &resultType,
        const core::QueryConfig & /*config*/)
        -> std::unique_ptr<exec::Aggregate> {
                POLLUX_CHECK_EQ(
                    argTypes.size(), 1, "{}: unexpected number of arguments", name);

                const auto keyType = resultType->childAt(0);

                if (keyType->providesCustomComparison()) {
                    return POLLUX_DYNAMIC_SCALAR_TYPE_DISPATCH(
                                createMapUnionAggregateWithCustomCompare,
                                keyType->kind(),
                                resultType);
                }

                return createMapAggregate<MapUnionAggregate>(resultType);
            },
            withCompanionFunctions,
            overwrite);
    }
} // namespace kumo::pollux::aggregate::prestosql
